import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False  # 正确显示负号

class RobotArmVisualizer:
    def __init__(self):
        # 初始参数设置
        self.theta1 = 45     # 关节1角度(度)
        self.theta2 = 45     # 关节2角度(度)
        # 严格保证 theta1 + theta2 + theta3 = 270度
        self.theta3 = 270 - self.theta1 - self.theta2  
        
        self.L1 = 60         # 连杆1长度
        self.L2 = 50         # 连杆2长度
        self.L3 = 40         # 连杆3长度
        
        # 创建图形和轴
        self.fig, self.ax = plt.subplots(figsize=(10, 10))
        plt.subplots_adjust(left=0.1, bottom=0.4)
        
        # 绘制运动范围散点
        self.workspace_points, = self.ax.plot([], [], 'g.', alpha=0.2, label='运动范围')
        
        # 绘制机械臂
        self.arm_line, = self.ax.plot([], [], 'o-', lw=5, color='blue', 
                                    markersize=10, markerfacecolor='red')
        self.end_effector_text = self.ax.text(0, 0, '', fontsize=12)
        
        # 标记目标工作空间
        self.ax.axvspan(60, 150, ymin=0, ymax=1, facecolor='yellow', alpha=0.2)
        self.ax.text(105, 20, '目标工作空间', ha='center', fontsize=10, 
                    bbox=dict(facecolor='yellow', alpha=0.5))
        
        # 绘制边界线
        self.ax.axvline(x=60, color='red', linestyle='--', alpha=0.7, label='X=60')
        self.ax.axvline(x=150, color='red', linestyle='--', alpha=0.7, label='X=150')
        self.ax.axhline(y=0, color='red', linestyle='--', alpha=0.7, label='Y=0')
        
        # 设置坐标轴
        self.ax.set_xlim(0, 250)
        self.ax.set_ylim(0, 250)
        self.ax.set_xlabel('X坐标')
        self.ax.set_ylabel('Y坐标')
        self.ax.set_title('三连杆机械臂运动范围（约束：θ1+θ2+θ3=270°）')
        self.ax.grid(True)
        self.ax.legend()
        self.ax.set_aspect('equal', adjustable='box')
        
        # 创建滑块
        slider_color = 'lightgoldenrodyellow'
        
        # 关节角度滑块
        ax_theta1 = plt.axes([0.1, 0.30, 0.8, 0.03], facecolor=slider_color)
        ax_theta2 = plt.axes([0.1, 0.25, 0.8, 0.03], facecolor=slider_color)
        
        # 显示角度约束信息
        self.constraint_text = self.fig.text(0.5, 0.22, 
                                           f'约束: θ1+θ2+θ3 = {self.theta1+self.theta2+self.theta3:.1f}°\nθ3 = {self.theta3:.1f}°', 
                                           ha='center', fontsize=10,
                                           bbox=dict(facecolor='lightblue', alpha=0.5))
        
        # 连杆长度滑块
        ax_L1 = plt.axes([0.1, 0.15, 0.8, 0.03], facecolor=slider_color)
        ax_L2 = plt.axes([0.1, 0.10, 0.8, 0.03], facecolor=slider_color)
        ax_L3 = plt.axes([0.1, 0.05, 0.8, 0.03], facecolor=slider_color)
        
        # 创建滑块控件
        self.slider_theta1 = Slider(ax_theta1, r'$\theta_1$ (度)', -90, 180, valinit=self.theta1)
        self.slider_theta2 = Slider(ax_theta2, r'$\theta_2$ (度)', -180, 180, valinit=self.theta2)
        self.slider_L1 = Slider(ax_L1, r'$L_1$', 30, 100, valinit=self.L1)
        self.slider_L2 = Slider(ax_L2, r'$L_2$', 30, 100, valinit=self.L2)
        self.slider_L3 = Slider(ax_L3, r'$L_3$', 20, 80, valinit=self.L3)
        
        # 注册更新函数
        self.slider_theta1.on_changed(self.update_arm)
        self.slider_theta2.on_changed(self.update_arm)
        self.slider_L1.on_changed(self.update_workspace)
        self.slider_L2.on_changed(self.update_workspace)
        self.slider_L3.on_changed(self.update_workspace)
        
        # 初始化显示
        self.update_workspace(None)
        self.update_arm(None)
        
    def calculate_end_position(self, theta1, theta2):
        """计算末端位置，严格应用角度约束"""
        # 强制应用约束：theta3 = 270 - theta1 - theta2
        theta3 = 270 - theta1 - theta2
        
        # 转换为弧度
        t1 = np.radians(theta1)
        t2 = np.radians(theta2)
        t3 = np.radians(theta3)
        
        # 计算各关节绝对角度
        angle1 = t1
        angle2 = t1 + t2
        angle3 = t1 + t2 + t3  # 此角度应为 270°(4.712弧度)，确保末端垂直
        
        # 计算末端位置
        x = self.L1 * np.cos(angle1) + self.L2 * np.cos(angle2) + self.L3 * np.cos(angle3)
        y = self.L1 * np.sin(angle1) + self.L2 * np.sin(angle2) + self.L3 * np.sin(angle3)
        
        return x, y
    
    def generate_workspace(self):
        """生成运动范围，遍历过程中严格执行角度约束"""
        # 角度遍历范围（根据机械臂工作特性优化）
        theta1_range = np.arange(-90, 181, 3)  # 步长3°，提高精度
        theta2_range = np.arange(-180, 181, 3)
        
        points = []
        for t1 in theta1_range:
            for t2 in theta2_range:
                # 每次计算都严格应用约束
                x, y = self.calculate_end_position(t1, t2)
                # 只保留第一象限的点
                if x >= 0 and y >= 0:
                    points.append((x, y))
        
        return np.array(points)
    
    def update_workspace(self, val):
        """更新运动范围，基于当前连杆长度和角度约束"""
        self.L1 = self.slider_L1.val
        self.L2 = self.slider_L2.val
        self.L3 = self.slider_L3.val
        
        # 生成满足约束条件的运动范围
        workspace = self.generate_workspace()
        
        # 更新散点图
        if len(workspace) > 0:
            self.workspace_points.set_xdata(workspace[:, 0])
            self.workspace_points.set_ydata(workspace[:, 1])
        
        self.update_arm(None)
    
    def update_arm(self, val):
        """更新机械臂姿态，保持角度约束"""
        self.theta1 = self.slider_theta1.val
        self.theta2 = self.slider_theta2.val
        self.theta3 = 270 - self.theta1 - self.theta2  # 强制约束
        
        # 计算各关节位置
        positions = self.calculate_joint_positions()
        x = [p[0] for p in positions]
        y = [p[1] for p in positions]
        
        # 更新机械臂显示
        self.arm_line.set_xdata(x)
        self.arm_line.set_ydata(y)
        
        # 更新末端位置文本
        end_x, end_y = positions[-1]
        self.end_effector_text.set_text(f'末端位置: ({end_x:.2f}, {end_y:.2f})')
        self.end_effector_text.set_x(end_x)
        self.end_effector_text.set_y(end_y)
        
        # 更新约束信息显示
        self.constraint_text.set_text(
            f'约束: θ1+θ2+θ3 = {self.theta1+self.theta2+self.theta3:.1f}°\nθ3 = {self.theta3:.1f}°'
        )
        
        self.fig.canvas.draw_idle()
    
    def calculate_joint_positions(self):
        """计算各关节位置，应用角度约束"""
        t1 = np.radians(self.theta1)
        t2 = np.radians(self.theta2)
        t3 = np.radians(self.theta3)
        
        angle1 = t1
        angle2 = t1 + t2
        angle3 = t1 + t2 + t3  # 始终为270°(4.712弧度)
        
        # 各关节坐标
        x0, y0 = 0, 0
        x1 = x0 + self.L1 * np.cos(angle1)
        y1 = y0 + self.L1 * np.sin(angle1)
        x2 = x1 + self.L2 * np.cos(angle2)
        y2 = y1 + self.L2 * np.sin(angle2)
        x3 = x2 + self.L3 * np.cos(angle3)
        y3 = y2 + self.L3 * np.sin(angle3)
        
        return [(x0, y0), (x1, y1), (x2, y2), (x3, y3)]
    
    def show(self):
        plt.show()

if __name__ == "__main__":
    visualizer = RobotArmVisualizer()
    visualizer.show()
